import argparse
import json
import os

import ultralytics


# ============================Tip==============================
# 1.需要安装对应的python3.8级以上，并且安装yolo环境
# 2.脚本返回的是一个PredictResult实例，包含了所有预测信息和图片的保存路径
# =============================================================

class PredictInfo:
    # PredictInfo类用于存储单个预测结果的信息
    def __init__(self, class_id, name, confidence):
        self.class_id = class_id
        self.name = name
        self.confidence = confidence

    def __repr__(self):
        # 修改此方法以打印对象时显示详细信息
        return f"PredictInfo(class_id={self.class_id}, name='{self.name}', confidence={self.confidence})"


class PredictResult:
    # PredictResult类用于存储所有预测结果信息以及图片的保存路径
    def __init__(self, save_path):
        self.save_path = save_path
        self.predict_infos = []

    def add_predict_info(self, predict_info):
        self.predict_infos.append(predict_info)

    def __repr__(self):
        # 修改此方法以打印对象时显示详细信息
        return f"PredictResult(save_dir='{self.save_path}', predict_infos={self.predict_infos})"


def serialize_predict_result(predict_result):
    """将 PredictResult 对象序列化为 JSON 字符串"""
    return json.dumps({
        'save_path': predict_result.save_path,
        'predict_infos': [
            {'class_id': info.class_id, 'name': info.name, 'confidence': info.confidence}
            for info in predict_result.predict_infos
        ]
    }, indent=4)


if __name__ == "__main__":
    # 图片检测
    save_dir = os.path.join(os.getcwd(), "result", "run")
    parser = argparse.ArgumentParser()
    mode_path = os.path.join(os.getcwd(), "../../model/bird/last.pt")
    parser.add_argument('--model', type=str, default=mode_path)
    parser.add_argument('--path', type=str, default="E:\\ThisPC\\Pictures\\bird\\赤膀鸭.jpg")
    parser.add_argument('--save_dir', type=str, default=save_dir)
    parser.add_argument('--conf', type=float, default=0.50)
    args = parser.parse_args()

    model = ultralytics.YOLO(args.model)

    # source后为要预测的图片数据集的的路径
    # save=True为保存预测结果
    # save_conf=True为保存坐标信息
    # save_txt=True为保存txt结果，但是yolov8本身当图片中预测不到异物时，不产生txt文件
    # name为路径
    result = model.predict(
        source=args.path,
        save=True,
        save_conf=True,
        save_txt=False,
        name=args.save_dir,
        conf=args.conf,
        verbose=False
    )

    # 创建PredictResult实例，并添加保存路径
    save_path = os.path.join(result[0].save_dir, os.path.basename(args.path))
    predict_result = PredictResult(save_path)

    # 获取检测结果对象
    detections = result[0].boxes  # 假设xyxy包含了检框的信息
    # 输出每个检测框的类别信息
    if len(detections) > 0:
        # 获取类别信息字典
        classDict = result[0].names
        for detection in detections:
            class_id = detection.cls  # 获取类别标签
            class_conf = detection.conf  # 获取类别置信度
            # 获取识别类别
            # print(int(class_id[0]))
            # 根据字典获取识别的类别名称
            # print(classDict[int(class_id[0])])
            # 获取识别类别置信度
            # print(float(class_conf[0]))
            # 创建Predict实例并添加到结果中
            predict_info = PredictInfo(int(class_id[0]), classDict[int(class_id[0])], float(class_conf[0]))
            predict_result.add_predict_info(predict_info)
    # 将PredictResult实例序列化为JSON并打印
    print(serialize_predict_result(predict_result))
